import { WebSocketServer } from 'ws';

export function setupWebSocket(server) {
  try {
    const wss = new WebSocketServer({
      server,
      path: '/ws', // 指定 WebSocket 路径
      clientTracking: true,
      perMessageDeflate: {
        zlibDeflateOptions: {
          chunkSize: 1024,
          memLevel: 7,
          level: 3
        },
        zlibInflateOptions: {
          chunkSize: 10 * 1024
        }
      }
    });

    console.log('WebSocket server initialized');

    const clients = new Set();

    wss.on('listening', () => {
      console.log('WebSocket server is listening');
    });

    wss.on('connection', (ws, req) => {
      console.log('New WebSocket connection from:', req.socket.remoteAddress);

      clients.add(ws);

      // 设置心跳检测
      ws.isAlive = true;
      ws.on('pong', () => {
        ws.isAlive = true;
      });

      // 发送欢迎消息
      try {
        ws.send(JSON.stringify({
          type: 'welcome',
          message: 'Connected to WebSocket server'
        }));
      } catch (err) {
        console.error('Error sending welcome message:', err);
      }

      ws.on('message', (message) => {
        try {
          console.log('Received:', message.toString());
          // 处理接收到的消息
        } catch (err) {
          console.error('Error processing message:', err);
        }
      });

      ws.on('close', () => {
        console.log('Client disconnected');
        clients.delete(ws);
      });

      ws.on('error', (error) => {
        console.error('WebSocket client error:', error);
        clients.delete(ws);
      });
    });

    wss.on('error', (error) => {
      console.error('WebSocket server error:', error);
    });

    // 心跳检测定时器
    const heartbeatInterval = setInterval(() => {
      wss.clients.forEach((ws) => {
        if (ws.isAlive === false) {
          console.log('Terminating inactive client');
          return ws.terminate();
        }

        ws.isAlive = false;
        ws.ping();
      });
    }, 30000);

    wss.on('close', () => {
      console.log('WebSocket server closing');
      clearInterval(heartbeatInterval);
    });

    // 广播函数
    function broadcast(message) {
      const messageStr = typeof message === 'string' ? message : JSON.stringify(message);
      clients.forEach(client => {
        if (client.readyState === WebSocketServer.OPEN) {
          try {
            client.send(messageStr);
          } catch (err) {
            console.error('Error broadcasting to client:', err);
          }
        }
      });
    }

    return broadcast;
  } catch (err) {
    console.error('Error setting up WebSocket server:', err);
    throw err;
  }
}